import os
import numpy as np
from PIL import Image, UnidentifiedImageError

# 文件夹路径，包含所有图片文件
data_path = 'data_cat_dog'

total_pixels = 0
# 处理 RGB 图片
sum_normalized_pixel_values = np.zeros(3)

# 遍历文件夹中的图片文件
for root, dirs, files in os.walk(data_path):
    for file in files:
        if file.endswith(('.jpg', '.jpeg', '.png', 'bmp')):
            path = os.path.join(root, file)
            try:
                with Image.open(path) as image:
                    # 确保图像是 RGB 格式
                    if image.mode != "RGB":
                        image = image.convert("RGB")

                    # 转为 numpy 数组并归一化
                    image_array = np.array(image, dtype=np.float32) / 255.0

                    # 累积归一化后的像素值和像素数量
                    total_pixels += image_array.shape[0] * image_array.shape[1]  # H x W
                    sum_normalized_pixel_values += np.sum(image_array, axis=(0, 1))
            except UnidentifiedImageError:
                print(f"Skipping unidentified image: {path}")
            except Exception as e:
                print(f"Error processing {path}: {e}")
            
mean = sum_normalized_pixel_values / total_pixels
sum_squared_diff = np.zeros(3)

for root, dirs, files in os.walk(data_path):
    for file in files:
        if file.endswith(('.jpg', '.jpeg', '.png', 'bmp')):
            path = os.path.join(root, file)
            try:
                with Image.open(path) as image:
                    # 确保图像是 RGB 格式
                    if image.mode != "RGB":
                        image = image.convert("RGB")

                    # 转为 numpy 数组并归一化
                    image_array = np.array(image, dtype=np.float32) / 255.0

                    # 计算每个像素的平方差并累加
                    diff = (image_array - mean) ** 2
                    sum_squared_diff += np.sum(diff, axis=(0, 1))
            except UnidentifiedImageError:
                print(f"Skipping unidentified image: {path}")
            except Exception as e:
                print(f"Error processing {path}: {e}")

variance = sum_squared_diff / total_pixels

print("Mean:", mean)
print("Variance:", variance)
